structure COMPOSITION = struct
open REDUCE

(* =============================================================================================
 * Alpha-renames a type expression: creates a new context with fresh bindings and propagates it.
 * alphaRename: ty -> ty
 * ============================================================================================= *)
fun alphaRename ty1
      = applySubst (foldTy (fn (TyVar x, e) => solve [(TyVar x, nextVar() )] e
                             | (_,       e) => e
                           )
                           (ty1,initialEnv)
                   )
                   ty1

(* =============================================================================================
 * Converts type expressions into normal forms.
 * normalize: ctx -> ty -> ty
 * ============================================================================================= *)
fun (* if choices on the lhs and the rhs are the same, pick one choice for both lhs and rhs *)
    normalize _ (r as TyRule(TyTerm(t1,[z1 as TySum(a,b)]), TyTerm(t2,[z2 as TySum _])))
      = if t1 = t2 andalso z1 = z2
        then TyRule(TyTerm(t1,[a]), TyTerm(t2,[a]))
        else r

    (* map each choice on the lhs of a rule to the corresponding choice on the rhs *)
  | normalize env (TyRule(lhs as TyTerm(t,[TySum(a,b)]), TySum(x,y)))
      = let
          val x1 = normalize env (TyRule (TyTerm(t,[a]),x))
          val y1 = normalize env (TyRule (TyTerm(t,[b]),y))
        in
          TySum(x1,y1)
        end

    (* flatten rules by lifting sums from a rule's rhs to the outside of the rule:
     *   e.g. a -> (x + y) => a -> x + a -> y *)
  | normalize env (r as TyRule(a as TyTerm _, TySum(b,c)))
      = let
          (*val _ = fine("normalize: " ^ toString r) 1*)
          val ab = normalize env (TyRule (a,b))
          val ac = normalize env (alphaRename (TyRule (a,c)))
          val tyOut
                = (* if ac is an ID, then skip it, because ab is already a rule (ID is built into rules) *)
                  case ac of
                    TyRule(x,y) => if x = y then ab
                                   else TySum(ab,ac)
                  | _ => TySum(ab,ac)
          (*val _ = fine("normalize: result " ^ toString tyOut) ~1*)
        in
          tyOut
        end
  | normalize env (TyRule(lhs as TyTerm _, TyTerm(t2,[TySum(x,y)])) )
      = let
          val x1 = normalize env (TyRule (lhs, TyTerm(t2,[x])) )
          val y1 = normalize env (TyRule (alphaRename lhs, TyTerm(t2,[y])) )
        in
          TySum(x1,y1)
        end
  | normalize _ (r as TyRule(TyTerm _,_)) = r
  | normalize _ r = raise TypeError("normalize: Unexpected argument: " ^ toString r)

(* =============================================================================================
 * Checking conjunctive and sequential combinators:
 * checkStar: (ty * ty) -> context -> expr -> (ty * context)
 * checkSeq:  (ty * ty) -> context -> expr -> (ty * context)
 * ============================================================================================= *)
val isFoldS = ref false;

fun (* first-order composition: is well-typed if output of the first rule is type-unifiable with the premise of the second *)
    checkStar (TyRule(a,b as TyTerm _), r2 as TyRule(c,d as TyTerm _)) env expr
      = let
            val e2 = unify (b, [c], env, expr) (* on success, b and c are type-unifiable *)
            val ad = TyRule(applySubst e2 a, applySubst e2 d)
        in
            (ad, env)
            (* if hasVarPremise r2 orelse (isPrecise b andalso isPrecise c)*) (* guaranteed to apply to the output of r1 *)
            (* then (ad, e2)
               else (TySum(ad, TyRule(nextVar(), nextVar())), e2) *)
        end

    (* heterogeneous composition: first-order rule is on the left *)
  | checkStar (TyRule(a,b as TyTerm _), TyRule(c,d)) env expr
      = let
            val e2 = unify (b, [c], env, expr)
        in
            (TyRule(applySubst e2 a, TySum (applySubst e2 b,applySubst e2 d)), env)
            (* first- and higher-order results are combined with TySum *)
        end
    (* heterogeneous composition: first-order rule is on the right *)
  | checkStar (TyRule(a,b), TyRule(c,d as TyTerm _)) env expr
      = let
            val e2 = unify (a, [c], env, expr)
        in
            (TyRule(applySubst e2 a, TySum (applySubst e2 b,applySubst e2 d)), env)
        end

    (* higher-order composition: is well-typed if premises of the rules are type-unifiable *)
  | checkStar (TyRule(a,b), TyRule(c,d)) env expr
      = let
            val e2 = unify (a, [c], env, expr)
            val rhs = case b of
                        TyRule _  => TyList [b,d]
                      | TyList bs => TyList (bs@[d])
                      | arg       => raise TypeError("Match failure in checkStar on arg: " ^ toString arg)
            val (a2, rhs2) = (applySubst e2 a, mapTy (applySubst e2) rhs)
            val combTy = TyRule (a2, rhs2)
        in
            if not (!isFoldS)
            then (combTy, e2)
            else (isFoldS := false; (TyRule (a2, TyList[rhs2,alphaRename combTy]), e2) )
        end

  (* one or both of the operands are choice-based compositions *)
  | checkStar (a as TyRule _ , TySum(b,c)) env expr     (* a(b+c)=>ab+ac *)
      = let
            val (ab, _) = checkStar (a, b) env expr
            val (ac, _) = checkStar (alphaRename a, c) env expr
        in
            (TySum(ab,ac), env)
        end
  | checkStar (TySum(a,b), c as TyRule _ ) env expr
      = let
            val (ac, _) = checkStar (a, c) env expr
            val (bc, _) = checkStar (b, alphaRename c) env expr
        in
            (TySum(ac,bc), env)
        end
  | checkStar (TySum(a,b), TySum(c,d)) env expr         (* (a+b)(c+d)=>(ac+ad)+(bc+bd) *)
      = let
            val (ac, _) = checkStar (a            , c) env expr
            val (ad, _) = checkStar (alphaRename a, d) env expr
            val (bc, _) = checkStar (b            , alphaRename c) env expr
            val (bd, _) = checkStar (alphaRename b, alphaRename d) env expr
        in
            (TySum(TySum(ac,ad), TySum(bc,bd)), env)
        end

  (* composition with map iterators *)
  | checkStar (TyRule(a,b), m as TyMap _) env p
      = let
            val (ty2, env2) = reduce(m, b, env) p
        in
            (TyRule(a,ty2), env2)
        end
  | checkStar (m as TyMap _, r2 as TyRule(c as TyTerm(cRoot, cKids), d) ) env p
      = let
            val _ = debug ("checkStar: input: " ^ toString m ^ " | " ^ toString r2) 1
            (* since this is conjunction, incoming term's root must be the same as cRoot *)
            val a = TyTerm(cRoot, [nextVar()])
            val (b, env1) = reduce(m, a, env) p
            val a2 = applySubst env1 a
            val r1 = TyRule(a2, b)
            (* primitive types may produce identity rewrites; eliminate them since this is conjunctive composition *)
            val b2 = case b of
                       (TySum (x,y)) =>
                         if      a2 = x then y
                         else if a2 = y then x
                         else b
                     | _ => b
            (* check that the premise of r2 does not have terms applicable for the premise of r1 *)
            val _ = case a2 of
                      TyTerm(_,[TySome aKid]) =>
                        if List.exists (fn cKid => cKid = aKid) cKids
                        then raiseCompositionError r1 r2 p "application of the first operand produces a term incompatible with the second" 
                        else ()
                    | _ => ()
            (* infer structure of the incoming term through inverse rewriting of c *)
            val a3 = TyTerm(cRoot,map (fn cKid => case b2 of
                                                     TyTerm(_,[TySome bKid]) => if bKid = cKid (* found a match *)
                                                                              then case a2 of
                                                                                     (* replace existential quantifier by a concrete term *)
                                                                                     TyTerm(_,[TySome aKid]) => aKid
                                                                                   | _ => cKid
                                                                              else cKid
                                                   | _ => cKid) cKids)
            (* if the right operand's structure's unknown, the left operand's structure should remain as before *)
            val a4 = case cKids of
                       [TyVar _] => a2
                     | _         => a3
            val eF = unify (b2, [c], env, p)
            val tyF = TyRule(applySubst eF a4, applySubst eF d)
            val _ = debug ("checkStar: output: " ^ toString tyF) ~1
        in
            (tyF, eF)
        end

  (* one of the operands is of unknown type: perform type-inferencing *)
  | checkStar (x as TyVar _, y as TyRule(a,_)) e p = (y, unify(x, [TyRule(a,a)], e, p))
  | checkStar (x as TyRule(_,a), y as TyVar _) e p = (x, unify(y, [TyRule(a,a)], e, p))
  | checkStar (a as TyVar _, TySum(b,c)) e p
      = let
            val (ab, e1) = checkStar (a,b) e p
            val (ac, e2) = checkStar (a,c) e p
        in
            (TySum(ab,ac), intersect(e1, e2))
        end
  | checkStar (TySum(a,b), c as TyVar _) e p
      = let
            val (ac, e1) = checkStar (a,c) e p
            val (bc, e2) = checkStar (b,c) e p
        in
            (TySum(ac,bc), intersect(e1, e2))
        end

  | checkStar (TyInf, _) e p = (TyInf, e)
  | checkStar (_, TyInf) e p = (TyInf, e)
  
  (* otherwise raise a type error *)
  | checkStar (ty1,ty2) _ p = raiseCompositionError ty1 ty2 p "unexpected operands in conjunctive composition"

(* ============================================================================================= *)
fun (* first-order composition *)
    checkSeq (r1 as TyRule(a,b as TyTerm _), r2 as TyRule(c,d as TyTerm _)) env expr
      = let
            val _ = debug("checkSeq(first-order): " ^ toString r1 ^ " and " ^ toString r2) 1
            val (unifiable, e2) = (true, unify (c, [b], env, expr)) handle TypeError _ => (false, env)
            val ad = TyRule(applySubst e2 a, applySubst e2 d)
            val isSubtypeBC = subtype b c
            val isSubtypeCB = subtype c b
            val isCondBetter
                = not unifiable (* output of the first rule is not compatible with the input of the second *)
                  andalso not (isSubtypeBC orelse isSubtypeCB) (* both rules' premises are mutually-exclusive *)
            val (tyOut, envOut)
                   = if unifiable
                     then if hasVarPremise r2 orelse (isPrecise b andalso isPrecise c) (* guaranteed to apply to the output of r1 *)
                          then if ad = r2
                               then (ad, e2)
                               else (TySum(ad, (*alphaRename*) r2), e2)
                          else (TySum(ad, TySum(alphaRename r1, (*alphaRename*) r2)), e2)
                     else (* log an info notice: the seq combinator can be replaced by a cond combinator *)
                          let
                              (*val _ = if isCondBetter
                                        then print2("Info: Sequential composition can be replaced by conditional composition" ^ crlf ^
                                                    "  left operand: " ^ toString r1 ^ crlf ^
                                                    "  right operand: " ^ toString r2 ^ crlf ^
                                                    "  in expression: " ^ A.leavesToString expr false ^ crlf)
                                        else ()*)
                          in
                              (TySum(r1, r2), env) (* todo: check if r2 is reachable *)
                          end
            val _ = debug("checkSeq result: " ^ toString tyOut) ~1
        in
            (tyOut, envOut)
        end

    (* heterogeneous composition: first-order rule is on the left *)
  | checkSeq (r1 as TyRule(a,b as TyTerm _), r2 as TyRule(c,d)) env expr
      = let
            val (unifiable, e2) = (true, unify (c, [b], env, expr)) handle TypeError _ => (false, env)
            val combTy = TyRule(a, TySum(applySubst e2 b,applySubst e2 d))
        in
            if unifiable
            then if hasVarPremise r2 orelse (isPrecise b andalso isPrecise c) (* guaranteed to apply to the output of r1 *)
                 then (TySum(combTy, alphaRename r2), e2)
                 else (TySum(combTy, TySum(alphaRename r1, alphaRename r2)), e2)
            else (TySum(r1, r2), env)
        end
    (* heterogeneous composition: first-order rule is on the right *)
  | checkSeq (r1 as TyRule(a,b), r2 as TyRule(c,d as TyTerm _)) env expr
      = let
            val (unifiable, e2) = (true, unify (c, [a], env, expr)) handle TypeError _ => (false, env)
            val combTy = TyRule(applySubst e2 a, TySum(applySubst e2 b,applySubst e2 d))
            val r1Alpha = alphaRename r1
            val r2Alpha = alphaRename r2
        in
            if unifiable
            then if (isPrecise r1 andalso isPrecise r2) orelse
                    (hasVarPremise r1 andalso hasVarPremise r2) (* both rules are guaranteed to fire *)
                 then (combTy, e2)
                 else if hasVarPremise r1   (* r2 is prim or precise *)
                      then (TySum(r1Alpha, combTy), env)
                      else if hasVarPremise r2
                           then (TySum(combTy, r2Alpha), env)
                           else (TySum(combTy, TySum (r1Alpha, r2Alpha)), env)
            else (TySum (r1,r2), env)
        end

    (* higher-order composition *)
  | checkSeq (r1 as TyRule(a,b), r2 as TyRule(c,d)) env expr
      = let
            val (unifiable, e2) = (true, unify (c, [a], env, expr)) handle TypeError _ => (false, env)
            val rhs = case b of
                        TyRule _  => TyList [b,d]
                      | TyList bs => TyList (bs@[d])
                      | arg       => raise TypeError("Match failure in checkSeq on arg: " ^ toString arg)
            val (a2, rhs2) = (applySubst e2 a, mapTy (applySubst e2) rhs)
            val combTy = TyRule (a2, rhs2)
            val r1Alpha = alphaRename r1
            val r2Alpha = alphaRename r2

            val isHetero = !isFoldS
            val _ = isFoldS := false
            val combTyHetero = TyRule (a2, TyList[rhs2,alphaRename combTy])
            val r1Hetero = TyRule(a, TyList[b,r1Alpha])
            val r2Hetero = TyRule(c, TyList[d,r2Alpha])
        in
            if not isHetero
            then if unifiable
                 then if (isPrecise r1 andalso isPrecise r2) orelse
                         (hasVarPremise r1 andalso hasVarPremise r2) (* both rules are guaranteed to fire *)
                      then (combTy, e2)
                      else if hasVarPremise r1   (* r2 is prim or precise *)
                           then (TySum(r1Alpha, combTy), env)
                           else if hasVarPremise r2
                                then (TySum(combTy, r2Alpha), env)
                                else (TySum(combTy, TySum (r1Alpha, r2Alpha)), env)
                 else (TySum (r1,r2), env)
            else if unifiable
                 then if (isPrecise r1 andalso isPrecise r2) orelse
                         (hasVarPremise r1 andalso hasVarPremise r2) (* both rules are guaranteed to fire *)
                      then (combTyHetero, e2)
                      else if hasVarPremise r1   (* r2 is prim or precise *)
                           then (TySum(r1Hetero, combTyHetero), env)
                           else if hasVarPremise r2
                                then (TySum(combTyHetero, r2Hetero), env)
                                else (TySum(combTyHetero, TySum (r1Hetero, r2Hetero)), env)
                 else (TySum (r1Hetero,r2Hetero), env)
        end

  (* one or both of the operands are choice-based compositions *)
  | checkSeq (a as TyRule _ , TySum(b,c)) env expr      (* a(b+c)=>ab+ac *)
      = let
            val (ab, _) = checkSeq (a,b) env expr
            val (ac, _) = checkSeq (alphaRename a,c) env expr
            (* simplification: if a(b+c)=>(a+b)+(a+c), then a(b+c)=>a+(b+c) *)
            val f = case ab of
                      TySum(a1,b1) => a1 = a andalso b1 = b andalso
                                      (case #1(checkSeq (a,c) env expr) of
                                              TySum(a2,c1)  => a2 = a andalso c1 = c
                                            | _ => false)
                    | _ => false
            val tyOut = if f then TySum(a,TySum(b,c)) else TySum(ab,ac) 
        in
            (tyOut, env)
        end
  | checkSeq (TySum(a,b), c as TyRule _ ) env expr (* (a+b)c=>ac+bc *)
      = let
            val (ac, _) = checkSeq (a,c) env expr
            val (bc, _) = checkSeq (b,alphaRename c) env expr
            (* simplification: if (a+b)+c=>(a+c)+(b+c), then (a+b)c=>(a+b)+c *)
            val f = case ac of
                      TySum(a1,c1) => a1 = a andalso c1 = c andalso
                                      (case #1(checkSeq (b,c) env expr) of
                                              TySum(b1,c2)  => b1 = b andalso c2 = c
                                            | _ => false)
                    | _ => false
            val tyOut = if f then TySum(TySum(a,b),c) else TySum(ac,bc) 
        in
            (tyOut, env)
        end
  | checkSeq (TySum(a,b), TySum(c,d)) env expr          (* (a+b)(c+d)=(ac+ad)+(bc+bd) *)
      = let
            val (ac, _) = checkSeq (a,            c            ) env expr
            val (ad, _) = checkSeq (alphaRename a,d            ) env expr
            val (bc, _) = checkSeq (b,            alphaRename c) env expr
            val (bd, _) = checkSeq (alphaRename b,alphaRename d) env expr
        in
            (TySum(TySum(ac,ad), TySum(bc,bd)), env)
        end

  (* composition with map iterators *)
  | checkSeq (TyRule(a,b), m as TyMap s) env p
      = let
            val (ty2, _) = reduce(m, b, env) p
            val case1 = normalize env (TyRule(a,ty2))
            
            val c = TyTerm("'term",[nextVar()]) (* generic node with variable children *)
            fun reduceToList (s as TyMap(TySum(a,b))) t env p
              = let
                  val (illTyped1, rhs1) = (false, reduceToList (TyMap a) t env p) handle TypeError _ => (true, [(t, env)])
                  val (illTyped2, rhs2) = (false, reduceToList (TyMap b) t env p) handle TypeError _ => (true, [(t, env)])
                in
                  if illTyped1 andalso illTyped2 then raiseAppError s t p "unexpected arguments in strategy application"
                  else if illTyped1 then rhs2
                  else if illTyped2 then rhs1
                  else rhs1@rhs2
                end
              | reduceToList s t env p = [reduce(s,t,env) p]
            val tuples = reduceToList m c env p
            val cases = map (fn (d,ctx) => TyRule(applySubst ctx c, d) ) tuples
            val cases2 = map (normalize env) cases
            val case2 = foldl (fn (term,acc) => TySum(acc,term)) (hd cases2) (tl cases2)
        in
            (TySum(case1,case2), env)
        end
  | checkSeq (m as TyMap _, r2 as TyRule(c as TyTerm(lhs,lhsKids),d) ) env p
      = let
            val _ = debug ("checkSeq: analyzing " ^ toString m ^ " <; " ^ toString r2) 1
            (* case 1: only mapL applies *)
            val term1 = TyTerm("'term", [nextVar()])
            val (out1, ctx1) = reduce(m, term1, env) p
            val rule1 = TyRule(applySubst ctx1 term1, out1)
            val _ = debug ("checkSeq: if only the left operand applies: " ^ toString rule1) 0
            (* case 2: only r2 applies *)
            val _ = debug ("checkSeq: if only the right operand applies: " ^ toString r2) 0
            (* case 3: both operands apply *)
            val (rule3, _) = checkStar (m, r2) env p
            val _ = debug ("checkSeq: if both operands apply: " ^ toString rule3) 0
            val composition = TySum(rule3, TySum(rule1, r2))
            val _ = debug ("checkSeq: result is " ^ toString composition) ~1
        in
            (composition, env)
        end

  (* one of the operands is of unknown type: perform type-inferencing *)
  | checkSeq (x as TyVar _, y as TyRule(a,_)) e p = (y, unify(x, [TyRule(a,a)], e, p))
  | checkSeq (x as TyRule(_,a), y as TyVar _) e p = (x, unify(y, [TyRule(a,a)], e, p))
  | checkSeq (a as TyVar _, TySum(b,c)) e p
      = let
            val (ab, e1) = checkSeq (a,b) e p
            val (ac, e2) = checkSeq (a,c) e p
        in
            (TySum(ab,ac), intersect(e1, e2))
        end
  | checkSeq (TySum(a,b), c as TyVar _) e p
      = let
            val (ac, e1) = checkSeq (a,c) e p
            val (bc, e2) = checkSeq (b,c) e p
        in
            (TySum(ac,bc), intersect(e1, e2))
        end
  
  | checkSeq (TyInf, _) e p = (TyInf, e)
  | checkSeq (_, TyInf) e p = (TyInf, e)

  (* otherwise raise a type error *)
  | checkSeq (ty1,ty2) _ p = raiseCompositionError ty1 ty2 p "unexpected operands in sequential composition"

(* =================================================================================================================
 * Composes strategies x and y using combinator F.
 * createComp: string -> ty -> ty -> context -> expr -> ty
 * ================================================================================================================= *)
fun createComp F x y env p = case F of
                               "<+>" => (TySum (x,y), env)
                             | "<+"  => (TySum (x,y), env)
                             | "+>"  => (TySum (y,x), env)
                             | "<;"  => checkSeq (x,y) env p
                             | ";>"  => checkSeq (y,x) env p
                             | "<*"  => checkStar(x,y) env p
                             | _     => checkStar(y,x) env p

(* =================================================================================================================
 * Folds a list of dynamic rules into a single strategy.
 * foldRules: string -> expr -> (ty * context) -> (ty * context)
 * foldSRules: string -> expr -> (ty * context) -> (ty * context)
 * ================================================================================================================= *)
fun foldRules F _ (TyRule(a, b as TyRule _  ), env) = (TyRule (a, b), env) (* nothing to fold *)
  | foldRules F p (TyRule(a, TyList (x::xs) ), env)
    = let
          val (ty1,env1) = foldl (fn (z,acc) => createComp F (#1 acc) z (#2 acc) p) (x,env) xs
      in
          (TyRule (applySubst env1 a, ty1), env1)
      end
  | foldRules F p (TySum (a,b), env)
    = let
          val (tyA, _) = foldRules F p (a,env)
          val (tyB, _) = foldRules F p (b,env)
      in
          (TySum(tyA, tyB), env)
      end
  | foldRules _ _ (TyInf, env) = (TyInf, env)
  | foldRules _ _ (ty1,_) = raise TypeError("Match failure in foldRules on arg: " ^ toString ty1)

fun foldSRules F p (r1 as TyRule(a, b as TyRule _ ), env) = foldRules F p (TyRule(a, TyList [b, alphaRename r1]), env)
  | foldSRules F p (r1 as TyRule(a, TyList (x::xs)), env)
    = let
          val (ty1,env1) = case x of
                             TyList (d::ds) => foldl (fn (z,acc) => createComp F (#1 acc) z (#2 acc) p) (d,env) ds
                           | _              => (x,env) (* nothing to pre-fold *)
      in
          foldRules F p (TyRule(a, TyList (ty1::xs)), env1)
      end
  | foldSRules F p (TySum (a,b), env)
    = let
          val (tyA, _) = foldSRules F p (a,env)
          val (tyB, _) = foldSRules F p (b,env)
      in
          (TySum(tyA, tyB), env)
      end
  | foldSRules _ expr (ty1,_)
      = raise TypeError(ninfoToString (A.getExprInfo expr) ^
                        "Type error: foldS operator's argument is not a higher-order strategy" ^ crlf ^
                        "  strategy:      " ^ toString ty1 ^ crlf ^
                        "  in expression: " ^ A.leavesToString expr false)



end